package com.hsbc.databasesync.utils;

import com.sun.deploy.util.StringUtils;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVRecord;
import org.apache.commons.csv.CSVPrinter;
import org.apache.commons.csv.CSVParser;
import org.postgresql.copy.CopyManager;
import org.postgresql.core.BaseConnection;

import java.io.*;
import java.nio.charset.Charset;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.*;

public class PGCopyInUtils {
    /**
     * 将表中的数据导出到本地文件
     *
     * @param connection   连接
     * @param filePath     文件路径
     * @param tableOrQuery 表名 或者查询语句
     * @throws SQLException SQLException
     * @throws IOException  IOException
     */
    public static void copyToFile(Connection connection, String filePath, String tableOrQuery)
            throws SQLException, IOException {
        FileOutputStream fileOutputStream = null;
        try {
            CopyManager copyManager = new CopyManager((BaseConnection) connection);
            fileOutputStream = new FileOutputStream(filePath);
            String copyOut = "COPY " + tableOrQuery + " TO STDOUT  WITH csv header ";
            System.out.println(copyOut);
            final long line = copyManager.copyOut(copyOut, fileOutputStream);
            System.out.println(line);
        } finally {
            if (fileOutputStream != null) {
                try {
                    fileOutputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 将文件中的数据导入到数据库中
     *
     * @param connection 连接
     * @param filePath   文件路径
     * @param tableName  表名
     * @return long  导入的行数
     * @throws SQLException SQLException
     * @throws IOException  IOException
     */
    public static long copyFromFile(Connection connection, String filePath, String tableName)
            throws SQLException, IOException {
        FileInputStream fileInputStream = null;
        try {
            CopyManager copyManager = new CopyManager((BaseConnection) connection);
            fileInputStream = new FileInputStream(filePath);
            String copyIn = "COPY " + tableName + " FROM STDIN WITH csv header ";
            System.out.println(copyIn);
            return copyManager.copyIn(copyIn, fileInputStream);
        } finally {
            if (fileInputStream != null) {
                try {
                    fileInputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 修改csv文件的表头
     *
     * @param sourceFilePath：需要修改的csv文件路路径
     * @param destFilePath：生产的新的csv文件路径
     * @param newAndOldHeaderNameMap：新旧表头名 对应关系
     */
    public static void modifyCsvHeader(String sourceFilePath, String destFilePath, Map<String, String> newAndOldHeaderNameMap) {
        Reader fileReader = null;
        Writer fileWriter = null;
        Appendable out = null;
        try {
            fileReader = new FileReader(new File(sourceFilePath));
            fileWriter = new FileWriter(new File(destFilePath));
            CSVParser csvParser = CSVParser.parse(fileReader, CSVFormat.DEFAULT.withHeader());
            // 旧表csv表头列表
            List<String> oldHeaderNames = csvParser.getHeaderNames();
            // 新表csv表头列表
            List<String> newHeaderNames = new ArrayList<>();
            // 有效旧表头，数据表
            List<String> usefullOldHeaderNames = new ArrayList<>();

            for (String oldHeaderName : oldHeaderNames) {
                if (newAndOldHeaderNameMap.containsKey(oldHeaderName)) {
                    usefullOldHeaderNames.add(oldHeaderName);
                    newHeaderNames.add(newAndOldHeaderNameMap.get(oldHeaderName));
                }
            }
            String newHeaders = StringUtils.join(newHeaderNames, ",") + System.getProperty("line.separator");
            fileWriter.write(newHeaders);
            for (Iterator<CSVRecord> it = csvParser.iterator(); it.hasNext(); ) {
                CSVRecord record = it.next();
                List<String> msgList = new ArrayList<>();
                for (String headerName : usefullOldHeaderNames) {
                    msgList.add(record.get(headerName));
                }
                fileWriter.write(StringUtils.join(msgList, ",") + System.getProperty("line.separator"));
            }
            final CSVPrinter printer = CSVFormat.DEFAULT.print(fileWriter);
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (fileReader != null) {
                try {
                    fileReader.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (fileWriter != null) {
                try {
                    fileWriter.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

}
